package com.wutong.erb.util;

import cn.hutool.jwt.JWTException;
import com.wutong.erb.exception.JwtParseException;
import io.jsonwebtoken.Claims;
import io.jsonwebtoken.ExpiredJwtException;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.MalformedJwtException;
import io.jsonwebtoken.security.Keys;
import io.jsonwebtoken.security.SignatureException;

import javax.crypto.SecretKey;
import java.nio.charset.StandardCharsets;
import java.util.Date;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * Jwt工具类
 * @author wutong
 */
public class JwtUtil {
    private static final String SECRET_KEY = "daf66e01593f61a15b857cf433aae03a005812b31234e149036bcc8dee755dbb";

    // 预先生成密钥避免重复计算
    private static final SecretKey KEY = Keys.hmacShaKeyFor(SECRET_KEY.getBytes(StandardCharsets.UTF_8));

    /**
     * 生成token
     * 默认过期时间10分钟
     * @param subject 主题
     * @param claims 自定义claims
     * @return token
     */
    public static String generateToken(String subject, Map<String, Object> claims) {
        return generateToken(subject, claims, 600_000L);
    }

    /**
     * 生成token
     * @param subject 主题
     * @param claims 自定义claims
     * @param expiration 过期时间，单位ms
     * @return token
     */
    public static String generateToken(String subject, Map<String, Object> claims, Long expiration) {
        return Jwts.builder()
                .subject(subject)
                .claims(claims)
                .signWith(KEY)
                .expiration(new Date(System.currentTimeMillis() + expiration))
                .compact();
    }

    /**
     * 生成token
     * @param subject 主题
     * @param claims 自定义claims
     * @param expiration 过期时间，单位秒
     * @param timeUnit 时间单位
     * @return token
     */
    public static String generateToken(String subject, Map<String, Object> claims, int expiration, TimeUnit timeUnit) {
        return Jwts.builder()
                .subject(subject)
                .claims(claims)
                .signWith(KEY)
                .expiration(new Date(System.currentTimeMillis() + timeUnit.toMillis(expiration)))
                .compact();
    }

    public static boolean validateToken(String token) {
        try {
            Jwts.parser()
                    .verifyWith(KEY)
                    .build()
                    .parseSignedClaims(token)
                    .getPayload();
            return true;
        } catch (ExpiredJwtException | IllegalArgumentException | SignatureException | MalformedJwtException e) {
            return false;
        }
    }

    /**
     * 解析token
     * @param token token
     * @return claims
     */
    public static Claims parseToken(String token) {
        return Jwts.parser()
                .verifyWith(KEY)
                .build()
                .parseSignedClaims(token)
                .getPayload();
    }

    public interface TokenClaimsExceptionHandler<R> {
        /**
         * 正常处理
         *
         * @param claims 自定义claims
         */
        R normalHandler(Claims claims) throws Exception;

        /**
         * 处理token过期异常
         */
        R handleExpiredJwtException() throws Exception;

        /**
         * 处理token解析异常
         */
        R handleIllegalArgumentException() throws Exception;

        /**
         * 处理签名异常
         */
        R handleSignatureException() throws Exception;

        /**
         * 处理token格式异常
         */
        R handleMalformedJwtException() throws Exception;

        /**
         * 处理其他异常
         * @param e 异常
         */
        R handleOtherException(Exception e) throws Exception;
    }

    public interface TokenClaimsExceptionHandlerByDefault<R> extends TokenClaimsExceptionHandler<R> {
        /**
         * 正常处理
         *
         * @param claims 自定义claims
         */
        @Override
        R normalHandler(Claims claims) throws Exception;

        /**
         * 处理token过期异常 | 处理token解析异常 | 处理签名异常 | 处理token格式异常
         * @return R
         * @throws Exception 异常
         */
        R defaultHandler() throws Exception;

        @Override
        default R handleExpiredJwtException() throws Exception {
            return defaultHandler();
        }

        @Override
        default R handleIllegalArgumentException() throws Exception {
            return defaultHandler();
        }

        @Override
        default R handleSignatureException() throws Exception {
            return defaultHandler();
        }

        @Override
        default R handleMalformedJwtException() throws Exception {
            return defaultHandler();
        }

        /**
         * 处理其他异常
         *
         * @param e 异常
         */
        @Override
        R handleOtherException(Exception e) throws Exception;
    }

    public static <R> R getTokenClaims(String token, TokenClaimsExceptionHandler<R> handler) {
        Claims claims;
        try {
            claims = parseToken(token);
            return handler.normalHandler(claims);
        } catch (ExpiredJwtException e) {
            try {
                return handler.handleExpiredJwtException();
            } catch (Exception ex) {
                throw new JwtParseException();
            }
        } catch (IllegalArgumentException e) {
            try {
                return handler.handleIllegalArgumentException();
            } catch (Exception ex) {
                throw new JwtParseException();
            }
        } catch (SignatureException e) {
            try {
                return handler.handleSignatureException();
            } catch (Exception ex) {
                throw new JwtParseException();
            }
        } catch (MalformedJwtException e) {
            try {
                return handler.handleMalformedJwtException();
            } catch (Exception ex) {
                throw new JwtParseException();
            }
        } catch (Exception e) {
            try {
                return handler.handleOtherException(e);
            } catch (Exception ex) {
                throw new JwtParseException();
            }
        }
    }

    public static <R> R getTokenClaims(String token, TokenClaimsExceptionHandlerByDefault<R> handler) {
        Claims claims;
        try {
            claims = parseToken(token);
            return handler.normalHandler(claims);
        } catch (ExpiredJwtException | IllegalArgumentException | SignatureException | MalformedJwtException e) {
            try {
                return handler.defaultHandler();
            } catch (Exception ex) {
                throw new JwtParseException();
            }
        } catch (Exception e) {
            try {
                return handler.handleOtherException(e);
            } catch (Exception ex) {
                throw new JwtParseException();
            }
        }
    }

    public static long getExpiration(String token) {
        return getTokenClaims(token, new TokenClaimsExceptionHandlerByDefault<>() {
            @Override
            public Long normalHandler(Claims claims) {
                return claims.getExpiration().getTime();
            }

            @Override
            public Long defaultHandler() {
                return 0L;
            }

            @Override
            public Long handleOtherException(Exception e) {
                return 0L;
            }
        });
    }
}
